from Interpolation import Interpolation, run
from numpy import linalg, mat, array

class Spline(Interpolation):
    def __init__(self):
        super().__init__()
        self.hs = []

    def addKey(self, x, y):
        super().addKey(x, y)
        if len(self.keys)<2: return
        self.hs = [pj.x-pi.x for pi, pj in zip(self.keys[:-1], self.keys[1:])]
        mus = [hi/(hi+hj) for hi, hj in zip(self.hs[:-1], self.hs[1:])]
        las = [1-mu for mu in mus]
        ds = [6*((p2.y-p1.y)/(p2.x-p1.x)-(p1.y-p0.y)/(p1.x-p0.x))/(p2.x-p0.x) for p0, p1, p2 in zip(self.keys[:-2], self.keys[1:-1], self.keys[2:])]
        las = [0]+las+[None]
        mus = [None]+mus+[0]
        ds = [0]+ds+[0]
        d = array(ds)
        A = []
        for i in range(len(self.keys)):
            row = [0]*len(self.keys)
            row[i] = 2
            if i!=0: row[i-1]=mus[i]
            if i!=len(self.keys)-1: row[i+1]=las[i]
            A.append(row)
        ms = linalg.solve(A, d).tolist()
        def compute(x):
            # binary search
            if x<self.keys[0].x or x>self.keys[-1].x: return 0
            begin = 0
            end = len(self.keys)-2
            # to debug
            while begin<end:
                # print('x', x, 'begin', begin, 'end', end)
                mid = (begin+end+1)//2
                if x>=self.keys[mid].x: begin=mid
                else: end=mid-1
            j=begin
            factor1 = ((self.keys[j+1].x - x)**3) * ms[j] / (6*self.hs[j])
            factor2 = ((x - self.keys[j].x)**3) * ms[j+1] / (6*self.hs[j])
            factor3 = (self.keys[j].y - ms[j]*(self.hs[j]**2)/6)*(self.keys[j+1].x-x)/self.hs[j]
            factor4 = (self.keys[j+1].y - ms[j+1]*(self.hs[j]**2)/6)*(x-self.keys[j].x)/self.hs[j]
            return factor1+factor2+factor3+factor4
        self.compute = compute

run(Spline())